import jieba
import pandas as pd
from jieba.analyse import set_stop_words, extract_tags
from utils import cut_sent, cutWord, remStopWord, decList


# 读取文件
def readCsv(url):
    data = pd.read_csv(url, names=['txt'], sep='aaa', encoding="GBK", engine='python')
    return data


def remove_one(row):
    return row[:1]


def find_index(row):
    return row.find("回 ")


# 章节判断
def jucChapter(data):
    chapter = 0
    data['head'] = data.txt.apply(remove_one)
    data['mid'] = data.txt.apply(find_index)
    data['len'] = data.txt.apply(len)
    for i in range(len(data)):
        if data['head'][i] == "第" and data['mid'][i] > 0 and data['len'][i] < 30:
            chapter += 1
        if chapter >= 40 and data['txt'][i] == "附录一：成吉思汗家族":
            chapter = 0
        data.loc[i, 'chap'] = chapter

    # 删除临时变量
    del data['head']
    del data['mid']
    del data['len']
    return data


# 给特定的章节添加段落号
def addPara(data, chap):
    the_chapter = data[data['chap'] == chap].copy()
    the_chapter.reset_index(drop=True, inplace=True)
    the_chapter['paraidx'] = the_chapter.index
    return the_chapter


# 分句并存到列表中
def saveList(the_chapter):
    chaplist = []
    for i in the_chapter['paraidx'][1:]:
        chaplist.append(cut_sent(the_chapter['txt'][i]))
    return chaplist


if __name__ == '__main__':
    chap = 1  # 章节数
    file = readCsv('./金庸-射雕英雄传txt精校版.txt')
    data = jucChapter(file)
    # 获取章节内容
    chapter = addPara(data, chap)
    chaplist = saveList(chapter)
    chapContent = ''.join(decList(chaplist))
    # 建立全文IDF词库
    allIdfWords = []
    for i in range(1, 40):
        chapter = addPara(data, i)
        chaplist = saveList(chapter)
        for para in chaplist:
            words = decList(cutWord(para, './【官方】金庸小说.txt'))
            remStopWords = remStopWord(words)
            allIdfWords += remStopWords
    for word in allIdfWords:
        with open('./射雕英雄传全文idf词库.txt', mode='a', encoding='utf-8') as f:
            f.write(word + '\n')
    # 提取关键词
    jieba.load_userdict('./射雕英雄传全文idf词库.txt')
    set_stop_words('./停用词.txt')
    mainWords = extract_tags(chapContent, withWeight=True)
    print(mainWords)